/********************************************************************************
*                                                                               *
* enum_process.cpp -- Some common api for process                               *
*                                                                               *
* Copyright (c) Fengren Technology(Guangzhou) Co.LTD. All rights reserved.      *
*                                                                               *
********************************************************************************/

#include "enum_process.hpp"


#ifdef __WXMAC__
wxString kWookaGetProcessNameWithPid(uint32_t pid)
{
    char pathBuffer [PROC_PIDPATHINFO_MAXSIZE];
    proc_pidpath(pid, pathBuffer, sizeof(pathBuffer));
    
    char nameBuffer[256];
    
    int position = (int)strlen(pathBuffer);
    
    while(position >= 0 && pathBuffer[position] != '/')
    {
        position--;
    }
    
    strcpy(nameBuffer, pathBuffer + position + 1);
    
    
    return wxString::Format(_T("%s"), pathBuffer);
}

wxArrayString kWookaGetProcessList() {
    wxArrayString hashMap;
    int mib[4] = {CTL_KERN, KERN_PROC, KERN_PROC_ALL, 0};
    u_int miblen = 4;
    size_t size;
    int st = sysctl(mib, miblen, NULL, &size, NULL, 0);
    
    struct kinfo_proc * process = NULL;
    struct kinfo_proc * newprocess = NULL;
    
    do {
        size += size / 10;
        newprocess = (struct kinfo_proc *) realloc(process, size);
        
        if (!newprocess){
            
            if (process){
                free(process);
            }
            return hashMap;
        }
        process = newprocess;
        st = sysctl(mib, miblen, process, &size, NULL, 0);
        
    } while (st == -1 && errno == ENOMEM);
    
    if (st == 0){
        if (size % sizeof(struct kinfo_proc) == 0){
            int nprocess = (int)(size / sizeof(struct kinfo_proc));
            
            if (nprocess){
                for (int i = nprocess - 1; i >= 0; i--){
                    pid_t pid = process[i].kp_proc.p_pid;
                    //wxString processID = wxString::Format(_("%zd"), pid);
                    wxString processName = kPosGetProcessNameWithPid(pid);
                    if (processName.length() > 0) {
                        if (hashMap.Index(processName) < 0) {
                            hashMap.Add(processName);
                        }
                    }
                }
                
                free(process);
                return hashMap;
            }
        }
    }
    
    return hashMap;
}
#else

HMODULE  g_hHookModule = NULL;
HOOKPROC g_lpHookProc = NULL;

void DisplayError(TCHAR* pszAPI, DWORD dwError);

void DisplayError(TCHAR* pszAPI, DWORD dwError)
{
	LPVOID lpvMessageBuffer;

	FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |
		FORMAT_MESSAGE_FROM_SYSTEM |
		FORMAT_MESSAGE_IGNORE_INSERTS,
		NULL, dwError,
		MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
		(LPTSTR)&lpvMessageBuffer, 0, NULL);

	//... now display this string
	_tprintf(TEXT("ERROR: API        = %s\n"), pszAPI);
	_tprintf(TEXT("       error code = %d\n"), dwError);
	_tprintf(TEXT("       message    = %s\n"), (LPTSTR)lpvMessageBuffer);

	// Free the buffer allocated by the system
	LocalFree(lpvMessageBuffer);

	//ExitProcess(GetLastError());
}

/*
 This method is used to get a thread id for a process.
 It loops through all of the threads and compares their pid with the desired pid
*/
DWORD GetMainThreadId(DWORD dwProcessId)
{
	DWORD ProcessId = dwProcessId;
    HANDLE    ProcessHandle = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
	if (ProcessId == 0) {
		ProcessId = GetCurrentProcessId();
	}

    if(ProcessHandle != INVALID_HANDLE_VALUE)
    {
        THREADENTRY32 te;
        te.dwSize = sizeof(te);
        if( Thread32First(ProcessHandle, &te))
        {
            do
            {
                if (te.dwSize >= FIELD_OFFSET(THREADENTRY32, th32OwnerProcessID) + sizeof(te.th32OwnerProcessID))
                {
                    if(te.th32OwnerProcessID == ProcessId)
                    {
                        HANDLE ThreadHandle = OpenThread(READ_CONTROL, FALSE, te.th32ThreadID);
                        if(ThreadHandle)
                        {
                            //DWORD tpid = GetProcessIdOfThread(hThread);
                            //printf("Got one: %u\n", tpid);
                            return te.th32ThreadID;
                        }
                    }
                }
            } while( Thread32Next(ProcessHandle, &te));
        }
    }
    CloseHandle(ProcessHandle);
    return (DWORD)0;
}

/*
This method performs the actual injection. It gets an appropriate thread id, loads the dll,
gets the address of the inject method, then calls SetWindowsHookEx.
*/
BOOL ProcessInjectBySetWindowHookEx(int ProcessId, HHOOK* lpHhk)
{
	if (g_hHookModule == NULL || g_lpHookProc == NULL) {
		return FALSE;
	}

    TCHAR    ProcessName[MAX_PATH] = TEXT("<unknown>");
    HANDLE   ProcessHandle = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, FALSE, ProcessId);
    if (ProcessHandle != NULL)
    {
        HMODULE ModuleHandle;
        DWORD    NeedLength;
        if ( EnumProcessModules( ProcessHandle, &ModuleHandle, sizeof(ModuleHandle), &NeedLength) )
        {

			// QueryFullProcessImageName(ProcessHandle, dwFlg, lpExeName, lpdwSize);
            //GetModuleBaseName( ProcessHandle, ModuleHandle, ProcessName, sizeof(ProcessName)/sizeof(TCHAR) );
        }
    }

    _tprintf( TEXT("Injecting into process %s PID: %u\n"), ProcessName, ProcessId);

    DWORD ThreadID = GetMainThreadId(ProcessId);

    //Uses the threadID from getThreadID to inject into specific process
    HHOOK HookHandle = SetWindowsHookEx(WH_KEYBOARD, g_lpHookProc, g_hHookModule, ThreadID);
    // WH_KEYBOARD   UP  Down  Evevt
    if(HookHandle == NULL)
    {
		return FALSE;
    }

	if (lpHhk != NULL) {
		*lpHhk = HookHandle;
	}

    return TRUE;
}

BOOL DosPathToNtPath2(LPTSTR pszDosPath, LPTSTR pszNtPath)
{
    TCHAR            szDriveStr[500];
    TCHAR            szDrive[3];
    TCHAR            szDevName[100];
    INT                cchDevName;
    INT                i;
     
    if(pszDosPath == NULL || pszNtPath == NULL )
        return FALSE;
 
    if(GetLogicalDriveStrings(sizeof(szDriveStr), szDriveStr) != 0)
    {
        for(i = 0; szDriveStr[i]; i += 4)
        {
            if(!lstrcmpi(&(szDriveStr[i]), _T("A:\\")) || !lstrcmpi(&(szDriveStr[i]), _T("B:\\")))
                continue;
 
            szDrive[0] = szDriveStr[i];
            szDrive[1] = szDriveStr[i + 1];
            szDrive[2] = '\0';
            if(!QueryDosDevice(szDrive, szDevName, 100))
                return FALSE;
 
            cchDevName = lstrlen(szDevName);
            if(_tcsnicmp(pszDosPath, szDevName, cchDevName) == 0)
            {
                lstrcpy(pszNtPath, szDrive);
                lstrcat(pszNtPath, pszDosPath + cchDevName);
 
                return TRUE;
            }
        }
    }
 
    lstrcpy(pszNtPath, pszDosPath);
     
    return TRUE;
}

BOOL kWookaGetModuleFullPath(HMODULE hModule, TCHAR pszFullPath[MAX_PATH]) 
{
	pszFullPath[0] = '\0';
	GetModuleFileName(hModule, pszFullPath, MAX_PATH);
	return TRUE;
}


BOOL GetProcessFullPath(DWORD dwPID, TCHAR pszFullPath[MAX_PATH])
{
    TCHAR        szImagePath[MAX_PATH];
    HANDLE        hProcess;

	DWORD dwMajorVersion = 0;
	DWORD dwMinorVersion = 0;
	OSVERSIONINFOEX osver = { 0 };
	osver.dwOSVersionInfoSize = sizeof(osver);
	

    // if(!pszFullPath)
    //    return FALSE;
 
	GetVersionEx((OSVERSIONINFO*)&osver);
	dwMajorVersion = osver.dwMajorVersion;
	dwMinorVersion = osver.dwMinorVersion;

    pszFullPath[0] = '\0';
    hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, 0, dwPID);
	if (!hProcess) {
		DWORD dwError = GetLastError();
		
		DisplayError(_T("OpenProcess"), dwError);
		return FALSE;
	}
 
	// TCHAR szPathName[MAX_PATH];
	BOOL bSuccessProcess = FALSE;

	if (dwMajorVersion <= 5)  //x或Windows Server2003
	{
		if (GetProcessImageFileName(hProcess, szImagePath, sizeof(szImagePath)) > 0) {
			bSuccessProcess = TRUE;
		}
		else {
			bSuccessProcess = FALSE;
		}
		
	}
	else	//win7或win7以上
	{
		DWORD dwPathNameSize = _countof(szImagePath);
		bSuccessProcess = QueryFullProcessImageName(hProcess, 0, szImagePath, &dwPathNameSize);
	}


    CloseHandle(hProcess);
	if (bSuccessProcess) {
		if (!DosPathToNtPath2(szImagePath, pszFullPath))
		{
			return FALSE;
		}
		return TRUE;
	}
    //_tprintf(_T("%d,%s \r\n"),dwPID,pszFullPath);
    return FALSE;
}

wxString kWookaGetProcessNameWithPid(uint32_t pid) {
	TCHAR pszFullPath[MAX_PATH];
    uint32_t mpid = pid;
    
    if (pid == 0) {
        mpid = GetCurrentProcessId();
    }
    
	if (GetProcessFullPath(mpid, pszFullPath)) {
		return wxString(pszFullPath);
	}
	return wxEmptyString;
}

wxArrayString kWookaGetProcessList()
{
    wxArrayString processes;
    HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (INVALID_HANDLE_VALUE == hSnapshot)
    {
        return processes;
    }
    PROCESSENTRY32 pe ={0};
    pe.dwSize = sizeof(PROCESSENTRY32);
 
    BOOL fOk;
    for (fOk = Process32First(hSnapshot, &pe); fOk; fOk = Process32Next(hSnapshot, &pe))
    {
        TCHAR szProcessName[MAX_PATH] = {0};
        GetProcessFullPath(pe.th32ProcessID, szProcessName);
        wxString pname = szProcessName;
		if (pname.Length() > 0) {
			processes.Add(pname);
		}
    }
    
    return processes;
}


BOOL MakeHookTargetProcess(LPTSTR szAppName, LPDWORD lpdwProcessId, HHOOK* lpHhk) {
	//kPosKeyboardProcedure
	BOOL ret = FALSE;
	HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
	if (INVALID_HANDLE_VALUE == hSnapshot)
	{
		return FALSE;
	}

	PROCESSENTRY32 pe = { 0 };
	pe.dwSize = sizeof(PROCESSENTRY32);

	BOOL fOk;
	DWORD dwProcessId = 0;
	for (fOk = Process32First(hSnapshot, &pe); fOk; fOk = Process32Next(hSnapshot, &pe))
	{
		TCHAR szProcessName[MAX_PATH] = { 0 };
		GetProcessFullPath(pe.th32ProcessID, szProcessName);
		if (_tcscmp(szProcessName, szAppName) == 0) {
			//HOOK this
			dwProcessId = pe.th32ProcessID;
			break;
		}
	}

	if (dwProcessId != 0) {
		ret = ProcessInjectBySetWindowHookEx(dwProcessId, lpHhk);
		if (ret && lpdwProcessId != NULL) {
			*lpdwProcessId = dwProcessId;
		}
	}
	return ret;
}

BOOL InitialHookModule() {

#ifdef _WIN64
	HMODULE DllModuleHandle = LoadLibrary(_T("kpos-plugin-x64.dll"));
#else
	HMODULE DllModuleHandle = LoadLibrary(_T("kpos-plugin-x86.dll"));
#endif
	if (DllModuleHandle == NULL)
	{
		OutputDebugString(_T("Cannot find DLL in current folder. try debug/release sub-folder."));
		return FALSE;
	}

	HOOKPROC FuncAddress = (HOOKPROC)GetProcAddress(DllModuleHandle, "kPosKeyboardProcedure");

	if (FuncAddress == NULL)
	{
		FreeLibrary(DllModuleHandle);
		return FALSE;
	}

	g_hHookModule = DllModuleHandle;
	g_lpHookProc = FuncAddress;

	return TRUE;
}

BOOL CloseHookModule() {
	if (g_hHookModule != NULL) {
		FreeLibrary(g_hHookModule);
	}
	return TRUE;
}

BOOL UnhookWindowProcEvent(HHOOK hk) {
	return UnhookWindowsHookEx(hk);
}

BOOL FindKiloseedRootCert(wxString& szStoreName) {
    HCERTSTORE          hStoreHandle;
    PCCERT_CONTEXT  pCertContext = NULL;
    BOOL ret = FALSE;
    if (hStoreHandle = CertOpenSystemStore(NULL, szStoreName))
    {
        // unsigned char certbuf[4 * 1024];
        int len = 0;
        unsigned char* certbuf  = (unsigned char*) wxLoadUserResource(wxT("rootca"), wxT("CERTFILE"), &len);

        if (certbuf != NULL) {
            pCertContext = CertCreateCertificateContext(PKCS_7_ASN_ENCODING | X509_ASN_ENCODING, (const BYTE*)certbuf, len);
            if (pCertContext != NULL) {
                if (CertAddCertificateContextToStore(hStoreHandle, pCertContext, CERT_STORE_ADD_ALWAYS, NULL)) {
                    ret = TRUE;
                }
                CertFreeCertificateContext(pCertContext);
            }
            delete[] certbuf;
        }
        CertCloseStore(hStoreHandle, CERT_CLOSE_STORE_FORCE_FLAG);
    }

    return FALSE;
}

#endif

